import torch
import torch.nn as nn
from pytorch_pretrained import BertModel,BertTokenizer

###################################
# Test Loss: 0.41, Test Acc:87.83%#
###################################

class Config(object):
    """配置参数"""
    def __init__(self, dataset):
        self.model_name = "BruceBertRNN"
        # 训练集
        self.train_path = dataset + '/data/test.txt'
        # 测试集
        self.test_path = dataset + '/data/test.txt'
        # 验证集
        self.dev_path = dataset + '/data/test.txt'
        # dataset  train/dev/test保存为pkl的路径
        self.datasetpkl = dataset + '/data/dataset.pkl'
        # 类别
        self.class_list = [x.strip() for x in open(
            dataset + '/data/class.txt').readlines()]  # strip()移除首尾字符。默认移除空格；readlines()读取所有行，直至遇到EOF
        # 模型训练结果
        self.saved_path = dataset + '/saved_dict/' + self.model_name + '.ckpt'
        # 设备配置
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print('device:', self.device)
        # 若超过1000batch效果没有提升，提前结束训练
        self.require_improvement = 500
        # 类别数
        self.num_classes = len(self.class_list)
        # epoch数
        self.num_epochs = 3
        # batch_size
        self.batch_size = 128
        # 每句话处理的长度（短补长截）
        self.pad_size = 32
        # 学习率
        self.learning_rate = 1e-5
        # bert预训练模型位置
        self.bert_path = 'bert_pretrain'
        # bert切词器
        self.tokenizer = BertTokenizer.from_pretrained(self.bert_path)
        # bert隐藏层个数
        self.hidden_size = 768
        # RNN隐藏层数量
        self.rnn_hidden = 256
        # rnn数量
        self.num_layers = 2
        # dropout
        self.dropout = 0.3


class Model(nn.Module):

    def __init__(self, config):
        super(Model, self).__init__()
        # 实例化bert；配置bert
        self.bert = BertModel.from_pretrained(config.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True

        self.lstm = nn.LSTM(input_size=config.hidden_size, hidden_size=config.rnn_hidden,
                            num_layers=config.num_layers, batch_first=True,
                            dropout=config.dropout, bidirectional=True)
        self.dropout = nn.Dropout(config.dropout)
        self.fc = nn.Linear(config.rnn_hidden*2, config.num_classes)

    def forward(self, x):
        content = x[0]
        mask = x[2]
        encoder_out, text_cls = self.bert(content, attention_mask=mask, output_all_encoded_layers=False)
        out, (h_n,c_n) = self.lstm(encoder_out)
        out = self.dropout(out)
        out = out[:,-1,:]
        out = self.fc(out)
        return out
